#!/usr/bin/python
# coding:utf-8

import argparse
import os
import os.path
import sys
import tempfile
from subprocess import check_call as call
from subprocess import check_output as callo

import yaml

HOME = os.path.dirname(os.path.abspath(sys.argv[0]))


def util_ansible_extra_vars(params):
    fd, name = tempfile.mkstemp('.yaml')
    os.close(fd)
    with open(name, 'w') as f:
        yaml.safe_dump(params, f)
    print(name)
    return '@' + name


def ansible_inv(config):
    fd, name = tempfile.mkstemp('.yaml')
    os.close(fd)
    with open(name, 'w') as f:

        # name-node
        f.writelines(['[name_node]\n'])
        for line in config['cluster']['name_node']:
            f.write(line + '\n')

        # format_name_node
        f.writelines(['[format_name_node]\n'])
        f.write(config['cluster']['name_node'][0] + '\n')

        # fetch_name_node
        f.writelines(['[fetch_name_node]\n'])
        for i, line in enumerate(config['cluster']['name_node']):
            if i > 0:
                f.write(line + '\n')

        # resource-manager
        f.writelines(['[resource_manager]\n'])
        for line in config['cluster']['resource_manager']:
            f.write(line + '\n')

        # hadoop-workers
        f.writelines(['[hadoop_workers]\n'])
        for line in config['cluster']['hadoop_workers']:
            f.write(line + '\n')

        # history-server
        f.writelines(['[history_server]\n'])
        for line in config['cluster']['history_server']:
            f.write(line + '\n')

        # zk-node
        f.writelines(['[zk_node]\n'])
        for index, line in enumerate(config['cluster']['zk_node']):
            f.write(line + ' myid=' + str(index) + '\n')

        # kafka_broker
        f.writelines(['[kafka_broker]\n'])
        for index, line in enumerate(config['cluster']['kafka_broker']):
            f.write(line + ' broker_id=' + str(index) + '\n')

        # docker_node
        f.writelines(['[docker_node]\n'])
        for index, line in enumerate(config['cluster']['docker_node']):
            f.write(line + '\n')

        all = set(
            config['cluster']['name_node'] + config['cluster']['name_node'] + config['cluster']['resource_manager'] +
            config['cluster']['hadoop_workers'] + config['cluster']['history_server'] + config['cluster']['zk_node'] +
            config['cluster']['kafka_broker'])
        f.writelines(['[all]\n'])
        for line in all:
            f.write(line + ' host_name=' + line + '\n')

    print(name)
    return name


def ansible_ssh_inv(config):
    fd, name = tempfile.mkstemp('.yaml')
    os.close(fd)
    with open(name, 'w') as f:
        # ssh_init
        f.writelines(['[ssh_init]\n'])
        for node in config['cluster']['all_node']:
            f.write("%s ansible_ssh_user=%s ansible_ssh_pass=%s ssh_inner_ip=%s\n"
                    % (node["host"],
                       config['cluster']['user']["name"],
                       config['cluster']['user']["password"],
                       node["inner_ip"]))
    print(name)
    return name


def ansible_ssh_inv_specific(config, host):
    fd, name = tempfile.mkstemp('.yaml')
    os.close(fd)
    with open(name, 'w') as f:
        # ssh_init
        f.writelines(['[ssh_init]\n'])
        for node in [host]:
            f.write("%s ansible_ssh_user=%s ansible_ssh_pass=%s ssh_inner_ip=%s\n"
                    % (node["host"],
                       config['cluster']['user']["name"],
                       config['cluster']['user']["password"],
                       node["inner_ip"]))
    print(name)
    return name


def cluster_ansible_vars(config):
    user = config['cluster']['user']

    zk_cluster_server = ""
    for i, node in enumerate(config['cluster']['zk_node']):
        zk_cluster_server += "server.%s=%s:2888:3888\n" % (i, node)

    etc_hosts = ""
    for node in config['cluster']["all_node"]:
        host = node['host']
        inner_ip = node['inner_ip']
        etc_hosts += "%s %s\n" % (inner_ip, host)

    vars = {
        'user_name': user['name'],
        'execute_user': user['name'],
        'user_password': user['password'],
        'zk_cluster_server': zk_cluster_server,
        'etc_hosts': etc_hosts
    }

    return util_ansible_extra_vars(vars)


def ssh_authorized_vars(config):

    host_list = []

    for node in config['cluster']["all_node"]:
        host = {
            "hostname": node['host'],
            "inner_ip": node['inner_ip']
        }
        host_list.append(host)

    vars = {
        'host_ip': host_list
    }

    return util_ansible_extra_vars(vars)


def server_prepare(config):
    v = cluster_ansible_vars(config)
    inv = ansible_inv(config)

    call(['ansible-playbook',
          '-i',
          inv,
          '-e',
          v,
          os.path.join(HOME, 'ansible', 'server_prepare.yml')])
    print("设置成功,需要重启服务器生效!")


def server_prepare_specific(host_name, config):
    v = cluster_ansible_vars(config)

    name = ansible_inv_specific(host_name)

    call(['ansible-playbook',
          '-i',
          name,
          '-e',
          v,
          os.path.join(HOME, 'ansible', 'server_prepare.yml')])
    print("设置成功,需要重启服务器生效!")


def ansible_inv_specific(host_name):
    fd, name = tempfile.mkstemp('.yaml')
    os.close(fd)
    with open(name, 'w') as f:
        node = [host_name]
        f.writelines(['[all]\n'])
        for line in node:
            f.write(line + ' host_name=' + line + '\n')
    print(name)
    return name


def install_jdk(config):
    v = cluster_ansible_vars(config)
    inv = ansible_inv(config)

    call(['ansible-playbook',
          '-i',
          inv,
          '-e',
          v,
          os.path.join(HOME, 'ansible', 'install_jdk.yml')])


def ssh_authorized(config):
    v = ssh_authorized_vars(config)
    inv = ansible_ssh_inv(config)
    all = set(
        config['cluster']['name_node'] + config['cluster']['name_node'] + config['cluster']['resource_manager'] +
        config['cluster']['hadoop_workers'] + config['cluster']['history_server'] + config['cluster']['zk_node'] +
        config['cluster']['kafka_broker'])
    do_ssh_authorize(config, inv, v, all)


def ssh_authorized_specific(host_name, ip, config):
    v = ssh_authorized_vars(config)
    inv = ansible_ssh_inv_specific(config, {"host": host_name, "inner_ip": ip})
    do_ssh_authorize(config, inv, v, [host_name])


def do_ssh_authorize(config, inv, v, all_node):
    call(['ansible-playbook',
          '-i',
          inv,
          '-e',
          v,
          os.path.join(HOME, 'ansible', 'ssh_authorized.yml')])
    with open(os.path.expandvars('$HOME') + "/.ssh/id_rsa.pub", 'r') as id_rsa_pub:
        id_rsa_pub_text = id_rsa_pub.read()
        for line in all_node:
            call([
                'ssh',
                'root@' + line,
                'echo \"' + id_rsa_pub_text + '\" >> /home/' + config['cluster']['user']['name'] + '/.ssh/authorized_keys'
            ])


def install_zookeeper(config):
    v = cluster_ansible_vars(config)
    inv = ansible_inv(config)

    call(['ansible-playbook',
          '-i',
          inv,
          '-e',
          v,
          os.path.join(HOME, 'ansible', 'install_zookeeper.yml')])


def install_hadoop(config):
    v = cluster_ansible_vars(config)
    inv = ansible_inv(config)

    var_replace = hadoop_cluster_vars(config)

    files = ['core-site.xml', 'hdfs-site.xml', 'mapred-site.xml', 'yarn-site.xml']
    for file_name in files:
        with open(os.path.join(HOME, 'conf', 'hadoop', file_name), 'r') as template_file:
            core_site_template = template_file.read()
            for key, value in var_replace.items():
                core_site_template = core_site_template.replace(key, value)
            with open(os.path.join(HOME, 'ansible', 'hadoop', file_name), 'w') as replace_file:
                replace_file.write(core_site_template)

    with open(os.path.join(HOME, 'ansible', 'hadoop', "workers"), 'w') as workers:
        for work in config['cluster']['hadoop_workers']:
            workers.writelines(work + "\n")

    call(['ansible-playbook',
          '-i',
          inv,
          '-e',
          v,
          os.path.join(HOME, 'ansible', 'install_hadoop.yml')])


def install_kafka(config):
    v = cluster_ansible_vars(config)
    inv = ansible_inv(config)

    var_zk_cluster = ""
    for i, node in enumerate(config['cluster']['zk_node']):
        var_zk_cluster += node + ":2181" + (',' if i < (len(config['cluster']['zk_node']) - 1) else '')
    var_zk_cluster += "/kafka"
    call(['ansible-playbook',
          '-i',
          inv,
          '-e',
          v,
          '-e',
          'zookeeper_connect=' + var_zk_cluster,
          os.path.join(HOME, 'ansible', 'install_kafka.yml')])


def hadoop_cluster_vars(config):
    var_user = config['cluster']['user']["name"]
    var_zk_cluster = ""
    for i, node in enumerate(config['cluster']['zk_node']):
        var_zk_cluster += node + ":2181" + (',' if i < (len(config['cluster']['zk_node']) - 1) else '')
    var_nn_comma = ""
    var_name_node_rpc = ""
    var_name_node_web = ""
    var_jn_url = "qjournal://"
    for i, node in enumerate(config['cluster']['name_node']):
        var_nn_comma += "nn" + str(i) + (',' if i < (len(config['cluster']['name_node']) - 1) else '')
        var_name_node_rpc += "\t<property>\n"
        var_name_node_rpc += "\t\t<name>dfs.namenode.rpc-address.mycluster.nn" + str(i) + "</name>\n"
        var_name_node_rpc += "\t\t<value>" + node + ":9820</value>\n"
        var_name_node_rpc += "\t</property>\n"

        var_name_node_web += "\t<property>\n"
        var_name_node_web += "\t\t<name>dfs.namenode.http-address.mycluster.nn" + str(i) + "</name>\n"
        var_name_node_web += "\t\t<value>" + node + ":9870</value>\n"
        var_name_node_web += "\t</property>\n"

        var_jn_url += node + ":8485" + (';' if i < (len(config['cluster']['name_node']) - 1) else '')
    var_jn_url += "/mycluster"
    var_history_server = config['cluster']['history_server'][0]
    var_rm_comma = ""
    var_rm_config = ""
    for i, node in enumerate(config['cluster']['resource_manager']):
        var_rm_comma += "rm" + str(i) + (',' if i < (len(config['cluster']['resource_manager']) - 1) else '')
        var_rm_config += "\t<property>\n"
        var_rm_config += "\t\t<name>yarn.resourcemanager.hostname.rm" + str(i) + "</name>\n"
        var_rm_config += "\t\t<value>" + node + "</value>\n"
        var_rm_config += "\t</property>\n"
        var_rm_config += "\t<property>\n"
        var_rm_config += "\t\t<name>yarn.resourcemanager.webapp.address.rm" + str(i) + "</name>\n"
        var_rm_config += "\t\t<value>" + node + ":8088</value>\n"
        var_rm_config += "\t</property>\n"
        var_rm_config += "\t<property>\n"
        var_rm_config += "\t\t<name>yarn.resourcemanager.address.rm" + str(i) + "</name>\n"
        var_rm_config += "\t\t<value>" + node + ":8032</value>\n"
        var_rm_config += "\t</property>\n"
        var_rm_config += "\t<property>\n"
        var_rm_config += "\t\t<name>yarn.resourcemanager.scheduler.address.rm" + str(i) + "</name>\n"
        var_rm_config += "\t\t<value>" + node + ":8030</value>\n"
        var_rm_config += "\t</property>\n"
        var_rm_config += "\t<property>\n"
        var_rm_config += "\t\t<name>yarn.resourcemanager.resource-tracker.address.rm" + str(i) + "</name>\n"
        var_rm_config += "\t\t<value>" + node + ":8031</value>\n"
        var_rm_config += "\t</property>\n"
    var_replace = {
        "var_user": var_user,
        "var_zk_cluster": var_zk_cluster,
        "var_nn_comma": var_nn_comma,
        "var_name_node_rpc": var_name_node_rpc,
        "var_name_node_web": var_name_node_web,
        "var_jn_url": var_jn_url,
        "var_history_server": var_history_server,
        "var_rm_comma": var_rm_comma,
        "var_rm_config": var_rm_config
    }
    return var_replace


def install_docker(config):
    v = cluster_ansible_vars(config)
    inv = ansible_inv(config)

    call(['ansible-playbook',
          '-i',
          inv,
          '-e',
          v,
          os.path.join(HOME, 'ansible', 'install_docker.yml')])


def command_server_prepare(args, config):
    if args.host_name:
        server_prepare_specific(args.host_name[0], config)
    else:
        server_prepare(config)


def command_server_ssh_authorized(args, config):
    if args.host_name:
        ssh_authorized_specific(args.host_name[0], args.ip[0], config)
    else:
        ssh_authorized(config)


def command_install_jdk(args, config):
    install_jdk(config)


def command_install_zookeeper(args, config):
    install_zookeeper(config)


def command_install_hadoop(args, config):
    install_hadoop(config)


def command_install_kafka(args, config):
    install_kafka(config)


def command_install_docker(args, config):
    install_docker(config)



def parse_args():
    parser = argparse.ArgumentParser()

    subparsers = parser.add_subparsers(dest='command')

    #Server
    server_parser = subparsers.add_parser('server', help='服务器操作')
    server_subparsers = server_parser.add_subparsers(dest='command2')
    prepare_parser = server_subparsers.add_parser('prepare', help='初始化操作系统')
    prepare_parser.add_argument("--host_name", help="可选，初始化指定服务器", nargs=1)
    ssh_authorized_parser = server_subparsers.add_parser('ssh-authorized', help='ssh免密登录')
    ssh_authorized_parser.add_argument("--host_name", help="可选，初始化指定服务器", nargs=1)
    ssh_authorized_parser.add_argument("--ip", help="可选，服务器IP", nargs=1)
    server_subparsers.add_parser('prepare-rollback', help='回滚初始化操作系统(TODO)')

    #install
    install_parser = subparsers.add_parser('install', help='安装服务')
    install_subparsers = install_parser.add_subparsers(dest='command2')
    install_subparsers.add_parser('jdk', help='安装jdk')
    install_subparsers.add_parser('zookeeper', help='安装Zookeeper')
    install_subparsers.add_parser('hadoop', help='安装Hadoop')
    install_subparsers.add_parser('kafka', help='安装Kafka')
    install_subparsers.add_parser('docker', help='安装Docker & DockerCompose')
    install_subparsers.add_parser('mysql', help='安装Mysql(docker)')

    #uninstall
    uninstall_parser = subparsers.add_parser('uninstall', help='卸载服务(TODO)')
    uninstall_subparsers = uninstall_parser.add_subparsers(dest='command2')
    uninstall_subparsers.add_parser('jdk', help='卸载jdk')
    uninstall_subparsers.add_parser('zookeeper', help='卸载Zookeeper')
    uninstall_subparsers.add_parser('hadoop', help='卸载Hadoop')
    uninstall_subparsers.add_parser('kafka', help='卸载Kafka')

    argv = sys.argv
    return parser.parse_args(argv[1:])


def run_args():
    args = parse_args()
    print("运行参数:%s" % args)

    command = args.command.replace('-', '_')
    command2 = args.command2
    if command2:
        command2 = command2.replace('-', '_')

    if command2:
        func = eval('command_%s_%s' % (command, command2))
    else:
        func = eval('command_%s' % (args.command))

    with open('cluster.yml', 'rb') as f:
        config = yaml.safe_load(f)
    func(args, config)


def main():
    run_args()
    return
    # with open('cluster.yml', 'rb') as f:
    #     config = yaml.safe_load(f)
    #     ssh_authorized_specific("app-server-001", "192.168.0.109", config)



if __name__ == '__main__':
    main()
